import torch
import torch.nn as nn
from torch_scatter import scatter_mean, scatter_add
from einops import rearrange

import torch
import numpy as np
import torch.nn as nn
from timm.models.layers import trunc_normal_
from einops import rearrange, repeat
from torch_scatter import scatter_softmax, scatter
from torch.nn import functional as F
from torch_geometric.nn.pool import global_add_pool,global_mean_pool

ACTIVATION = {
    "gelu": nn.GELU,
    "tanh": nn.Tanh,
    "sigmoid": nn.Sigmoid,
    "relu": nn.ReLU,
    "leaky_relu": nn.LeakyReLU(0.1),
    "softplus": nn.Softplus,
    "ELU": nn.ELU,
    "silu": nn.SiLU,
}


class Proj2Slices(nn.Module):
    def __init__(self, dim, heads=8, dim_head=64, dropout=0.0, slice_num=64):
        super().__init__()
        inner_dim = dim_head * heads
        self.dim_head = dim_head
        self.heads = heads
        self.scale = dim_head**-0.5
        self.slice_num = slice_num
        self.softmax = nn.Softmax(dim=-1)
        self.dropout = nn.Dropout(dropout)

        self.temperature = nn.Parameter(torch.ones([1, heads, 1, 1]) * 0.5)
        self.graph_temperature = nn.Parameter(torch.ones([1, heads, 1]) * 0.5)

        self.in_project_x = nn.Linear(dim, inner_dim)
        self.in_project_fx = nn.Linear(dim, inner_dim)
        self.in_project_slice = nn.Linear(dim_head, slice_num)
        for l in [self.in_project_slice]:
            torch.nn.init.orthogonal_(l.weight)  # use a principled initialization
        self.to_q = nn.Linear(dim_head, dim_head, bias=False)
        self.to_k = nn.Linear(dim_head, dim_head, bias=False)
        self.to_v = nn.Linear(dim_head, dim_head, bias=False)
        self.to_out = nn.Sequential(nn.Linear(inner_dim, dim), nn.Dropout(dropout))

    def graph_forward(self, x, batch): # return B, Slice, dim
        # x: [B*N, C], batch: [B*N, 1] 
        total_nodes = x.size(0)
        batch_size = batch.max().item() + 1

        ### (1) Slice
        fx_mid = self.in_project_fx(x).view(
            total_nodes, self.heads, self.dim_head
        )  # [B*N, H, D]
        x_mid = self.in_project_x(x).view(
            total_nodes, self.heads, self.dim_head
        )  # [B*N, H, D]

        slice_weights = self.softmax(
            self.in_project_slice(x_mid) / self.graph_temperature
        )  # [B*N, H, G]
        slice_norm = scatter_add(slice_weights, batch, dim=0)  # [B, H, G]

        slice_token = scatter_add(
            fx_mid.unsqueeze(-2) * slice_weights.unsqueeze(-1), batch, dim=0
        )  # [B, H, G, C]
        slice_token = slice_token / (slice_norm.unsqueeze(-1) + 1e-5)

        ### (2) Attention among slice tokens
        q_slice_token = self.to_q(slice_token)
        k_slice_token = self.to_k(slice_token)
        v_slice_token = self.to_v(slice_token)
        dots = torch.matmul(q_slice_token, k_slice_token.transpose(-1, -2)) * self.scale
        attn = self.softmax(dots)
        # attn = self.dropout(attn)
        out_slice_token = torch.matmul(attn, v_slice_token)  # B H G D
        out_slice_token = out_slice_token.permute(0, 2, 1, 3).reshape(batch_size, self.slice_num, -1)
        out_slice_token = self.to_out(out_slice_token)

        return out_slice_token


class MLP(nn.Module):
    def __init__(self, n_input, n_hidden, n_output, n_layers=1, act="gelu", res=True):
        super(MLP, self).__init__()

        if act in ACTIVATION.keys():
            act = ACTIVATION[act]
        else:
            raise NotImplementedError
        self.n_input = n_input
        self.n_hidden = n_hidden
        self.n_output = n_output
        self.n_layers = n_layers
        self.res = res
        self.linear_pre = nn.Sequential(nn.Linear(n_input, n_hidden), act())
        self.linear_post = nn.Linear(n_hidden, n_output)
        self.linears = nn.ModuleList(
            [
                nn.Sequential(nn.Linear(n_hidden, n_hidden), act())
                for _ in range(n_layers)
            ]
        )

    def forward(self, x):
        x = self.linear_pre(x)
        for i in range(self.n_layers):
            if self.res:
                x = self.linears[i](x) + x
            else:
                x = self.linears[i](x)
        x = self.linear_post(x)
        return x


class FeatureMixer(nn.Module):
    """Transformer encoder block."""

    def __init__(
        self,
        num_heads: int,
        hidden_dim: int,
        dim_head: int,
        dropout: float,
        act="gelu",
        mlp_ratio=2,
        slice_num=32,
        ufeature_ref = 8,
        outdim = 1
    ):
        super().__init__()
        self.proj2slice = Proj2Slices(
            dim=hidden_dim, 
            heads=num_heads, 
            dim_head=dim_head, 
            slice_num=slice_num
            )
        self.ref = ufeature_ref
        self.mixer = nn.MultiheadAttention(
            embed_dim=hidden_dim, 
            num_heads=1
            )
        self.out_mlp = MLP(
            n_input=hidden_dim, 
            n_hidden=hidden_dim*mlp_ratio,
            n_output=outdim, 
            n_layers=1, 
            res=False
            )


    def forward(self, gfeatures, batch, ufeatures):
        
        gfeatures = self.proj2slice.graph_forward(gfeatures, batch) # B, patch, C
        
        output_size = [self.ref]*3
        ufeatures = F.interpolate(ufeatures, size=output_size, mode="nearest") # B C D H W
        B, C, D, H, W = ufeatures.shape
        ufeatures = ufeatures.reshape(B, C, -1).permute(0, 2, 1) # B, patch, C
        
        features = torch.cat([gfeatures, ufeatures], dim=1) # B, patch, C
        features = features.permute(1, 0, 2) # L, B, Eqkv
        
        out, _ = self.mixer.forward(features, features, features) # L, B, Eqkv
        out = torch.mean(out, dim=0)
        out = self.out_mlp(out)
        return out
    



if __name__ == '__main__':
    B, N = 3, 4
    dim = 128
    x = torch.randn(B, N, dim).cuda() 
    batch = torch.tensor([0, 0, 0, 0, 1, 1, 1, 1, 2, 2, 2, 2]).cuda()
    x_graph = x.reshape(-1, dim).cuda() 
    
    p2s = Proj2Slices(128,8,64,0,7).cuda()
    
    cd = p2s.graph_forward(x.reshape(3*4, -1), batch)
    print(cd.shape)